{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b3af62b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "db856a45a72d464da338b8c13a035997",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 13 files:   0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The config attributes {'scaling_factor': 0.18215} were passed to AutoencoderKL, but are not expected and will be ignored. Please verify your config.json configuration file.\n",
      "You have disabled the safety checker for <class 'diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline'> by passing `safety_checker=None`. Ensure that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered results in services or applications open to the public. Both the diffusers team and Hugging Face strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling it only for use-cases that involve analyzing network behavior or auditing its results. For more information, please have a look at https://github.com/huggingface/diffusers/pull/254 .\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "('cuda',\n",
       " PNDMScheduler {\n",
       "   \"_class_name\": \"PNDMScheduler\",\n",
       "   \"_diffusers_version\": \"0.12.1\",\n",
       "   \"beta_end\": 0.012,\n",
       "   \"beta_schedule\": \"scaled_linear\",\n",
       "   \"beta_start\": 0.00085,\n",
       "   \"clip_sample\": false,\n",
       "   \"num_train_timesteps\": 1000,\n",
       "   \"prediction_type\": \"epsilon\",\n",
       "   \"set_alpha_to_one\": false,\n",
       "   \"skip_prk_steps\": true,\n",
       "   \"steps_offset\": 1,\n",
       "   \"trained_betas\": null\n",
       " },\n",
       " CLIPTokenizer(name_or_path='/root/.cache/huggingface/diffusers/models--lansinuote--diffsion_from_scratch.params/snapshots/310d9345e14b3b625635041dd573676c008d83ea/tokenizer', vocab_size=49408, model_max_length=77, is_fast=False, padding_side='right', truncation_side='right', special_tokens={'bos_token': AddedToken(\"<|startoftext|>\", rstrip=False, lstrip=False, single_word=False, normalized=True), 'eos_token': AddedToken(\"<|endoftext|>\", rstrip=False, lstrip=False, single_word=False, normalized=True), 'unk_token': AddedToken(\"<|endoftext|>\", rstrip=False, lstrip=False, single_word=False, normalized=True), 'pad_token': '<|endoftext|>'}))"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from diffusers import DiffusionPipeline\n",
    "import torch\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "pipeline = DiffusionPipeline.from_pretrained(\n",
    "    'lansinuote/diffsion_from_scratch.params', safety_checker=None)\n",
    "\n",
    "scheduler = pipeline.scheduler\n",
    "tokenizer = pipeline.tokenizer\n",
    "\n",
    "del pipeline\n",
    "\n",
    "device, scheduler, tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fd207966",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration lansinuote--diffsion_from_scratch-34318fc75271f5a0\n",
      "Found cached dataset parquet (/root/.cache/huggingface/datasets/lansinuote___parquet/lansinuote--diffsion_from_scratch-34318fc75271f5a0/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec)\n",
      "Loading cached processed dataset at /root/.cache/huggingface/datasets/lansinuote___parquet/lansinuote--diffsion_from_scratch-34318fc75271f5a0/0.0.0/2a3b91fbd88a2c90d1dbbb32b460cf621d31bd5b05b934492fdef7d8d6f236ec/cache-314ecfdae0cf40d9.arrow\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(Dataset({\n",
       "     features: ['pixel_values', 'input_ids'],\n",
       "     num_rows: 833\n",
       " }),\n",
       " {'pixel_values': tensor([[[1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           ...,\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.]],\n",
       "  \n",
       "          [[1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           ...,\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.]],\n",
       "  \n",
       "          [[1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           ...,\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "           [1., 1., 1.,  ..., 1., 1., 1.]]]),\n",
       "  'input_ids': tensor([49406,   320,  3610,   539,   320,  1901,  9528,   593,   736,  3095,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "          49407, 49407, 49407, 49407, 49407, 49407, 49407])})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "import torchvision\n",
    "\n",
    "#加载数据集\n",
    "dataset = load_dataset(path='lansinuote/diffsion_from_scratch', split='train')\n",
    "\n",
    "#图像增强模块\n",
    "compose = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.Resize(\n",
    "        512, interpolation=torchvision.transforms.InterpolationMode.BILINEAR),\n",
    "    torchvision.transforms.CenterCrop(512),\n",
    "    #torchvision.transforms.RandomHorizontalFlip(),\n",
    "    torchvision.transforms.ToTensor(),\n",
    "    torchvision.transforms.Normalize([0.5], [0.5]),\n",
    "])\n",
    "\n",
    "\n",
    "def f(data):\n",
    "    #应用图像增强\n",
    "    pixel_values = [compose(i) for i in data['image']]\n",
    "\n",
    "    #文字编码\n",
    "    input_ids = tokenizer.batch_encode_plus(data['text'],\n",
    "                                            padding='max_length',\n",
    "                                            truncation=True,\n",
    "                                            max_length=77).input_ids\n",
    "\n",
    "    return {'pixel_values': pixel_values, 'input_ids': input_ids}\n",
    "\n",
    "\n",
    "dataset = dataset.map(f,\n",
    "                      batched=True,\n",
    "                      batch_size=100,\n",
    "                      num_proc=1,\n",
    "                      remove_columns=['image', 'text'])\n",
    "\n",
    "dataset.set_format(type='torch')\n",
    "\n",
    "dataset, dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "69854652",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(833,\n",
       " {'pixel_values': tensor([[[[1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            ...,\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.]],\n",
       "  \n",
       "           [[1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            ...,\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.]],\n",
       "  \n",
       "           [[1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            ...,\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.],\n",
       "            [1., 1., 1.,  ..., 1., 1., 1.]]]], device='cuda:0'),\n",
       "  'input_ids': tensor([[49406,   320,  3610,   539,   320,  7651,  4009,   530,  3360,   537,\n",
       "            5046, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407,\n",
       "           49407, 49407, 49407, 49407, 49407, 49407, 49407]], device='cuda:0')})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#定义loader\n",
    "def collate_fn(data):\n",
    "    pixel_values = [i['pixel_values'] for i in data]\n",
    "    input_ids = [i['input_ids'] for i in data]\n",
    "\n",
    "    pixel_values = torch.stack(pixel_values).to(device)\n",
    "    input_ids = torch.stack(input_ids).to(device)\n",
    "\n",
    "    return {'pixel_values': pixel_values, 'input_ids': input_ids}\n",
    "\n",
    "\n",
    "loader = torch.utils.data.DataLoader(dataset,\n",
    "                                     shuffle=True,\n",
    "                                     collate_fn=collate_fn,\n",
    "                                     batch_size=1)\n",
    "\n",
    "len(loader), next(iter(loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "217c0a9c",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The config attributes {'scaling_factor': 0.18215} were passed to AutoencoderKL, but are not expected and will be ignored. Please verify your config.json configuration file.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(AdamW (\n",
       " Parameter Group 0\n",
       "     amsgrad: False\n",
       "     betas: (0.9, 0.999)\n",
       "     capturable: False\n",
       "     eps: 1e-08\n",
       "     foreach: None\n",
       "     lr: 1e-05\n",
       "     maximize: False\n",
       "     weight_decay: 0.01\n",
       " ),\n",
       " MSELoss())"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#加载模型\n",
    "%run 1.encoder.ipynb\n",
    "%run 2.vae.ipynb\n",
    "%run 3.unet.ipynb\n",
    "\n",
    "#准备训练\n",
    "encoder.requires_grad_(False)\n",
    "vae.requires_grad_(False)\n",
    "unet.requires_grad_(True)\n",
    "\n",
    "encoder.eval()\n",
    "vae.eval()\n",
    "unet.train()\n",
    "\n",
    "encoder.to(device)\n",
    "vae.to(device)\n",
    "unet.to(device)\n",
    "\n",
    "optimizer = torch.optim.AdamW(unet.parameters(),\n",
    "                              lr=1e-5,\n",
    "                              betas=(0.9, 0.999),\n",
    "                              weight_decay=0.01,\n",
    "                              eps=1e-8)\n",
    "\n",
    "criterion = torch.nn.MSELoss()\n",
    "\n",
    "optimizer, criterion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a3f32e9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_loss(data):\n",
    "    with torch.no_grad():\n",
    "        #文字编码\n",
    "        #[1, 77] -> [1, 77, 768]\n",
    "        out_encoder = encoder(data['input_ids'])\n",
    "\n",
    "        #抽取图像特征图\n",
    "        #[1, 3, 512, 512] -> [1, 4, 64, 64]\n",
    "        out_vae = vae.encoder(data['pixel_values'])\n",
    "        out_vae = vae.sample(out_vae)\n",
    "\n",
    "        #0.18215 = vae.config.scaling_factor\n",
    "        out_vae = out_vae * 0.18215\n",
    "\n",
    "    #随机数,unet的计算目标\n",
    "    noise = torch.randn_like(out_vae)\n",
    "\n",
    "    #往特征图中添加噪声\n",
    "    #1000 = scheduler.num_train_timesteps\n",
    "    #1 = batch size\n",
    "    noise_step = torch.randint(0, 1000, (1, )).long().to(device)\n",
    "    out_vae_noise = scheduler.add_noise(out_vae, noise, noise_step)\n",
    "\n",
    "    #根据文字信息,把特征图中的噪声计算出来\n",
    "    out_unet = unet(out_vae=out_vae_noise,\n",
    "                    out_encoder=out_encoder,\n",
    "                    time=noise_step)\n",
    "\n",
    "    #计算mse loss\n",
    "    #[1, 4, 64, 64],[1, 4, 64, 64]\n",
    "    return criterion(out_unet, noise)\n",
    "\n",
    "\n",
    "# get_loss({\n",
    "#     'input_ids': torch.ones(1, 77, device=device).long(),\n",
    "#     'pixel_values': torch.randn(1, 3, 512, 512, device=device)\n",
    "# })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "84b88599",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 11.7118999005761\n",
      "10 105.27776907754014\n",
      "20 101.45478522218764\n",
      "30 97.96161541804031\n",
      "40 95.7652038520173\n",
      "50 92.64628775657911\n",
      "60 91.62508884524868\n",
      "70 88.90302349776903\n",
      "80 84.6358380591555\n",
      "90 82.70271758512536\n",
      "100 81.53195204613439\n",
      "110 76.3927595877758\n",
      "120 74.14106083381193\n",
      "130 71.42537906522921\n",
      "140 69.16221529991162\n",
      "150 65.47076485656726\n",
      "160 62.1360088881047\n",
      "170 60.89056803673884\n",
      "180 57.985315461344726\n",
      "190 54.73302427918679\n",
      "200 50.69724302080431\n",
      "210 48.59712202517403\n",
      "220 46.407517315681616\n",
      "230 44.99496047659704\n",
      "240 44.07751854383969\n",
      "250 39.62402399040002\n",
      "260 37.051896732489695\n",
      "270 36.89249631060375\n",
      "280 35.71413582353853\n",
      "290 33.45783720578038\n",
      "300 33.08240255239798\n",
      "310 30.282505852694158\n",
      "320 29.86848702972202\n",
      "330 29.363934024146147\n",
      "340 29.187604612583527\n",
      "350 27.543819716789585\n",
      "360 26.130621485815936\n",
      "370 25.465440133120865\n",
      "380 25.48384229660587\n",
      "390 24.789676978944044\n"
     ]
    }
   ],
   "source": [
    "def train():\n",
    "    loss_sum = 0\n",
    "    for epoch in range(400):\n",
    "        for i, data in enumerate(loader):\n",
    "            loss = get_loss(data) / 4\n",
    "            loss.backward()\n",
    "            loss_sum += loss.item()\n",
    "\n",
    "            if (epoch * len(loader) + i) % 4 == 0:\n",
    "                torch.nn.utils.clip_grad_norm_(unet.parameters(), 1.0)\n",
    "                optimizer.step()\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "        if epoch % 10 == 0:\n",
    "            print(epoch, loss_sum)\n",
    "            loss_sum = 0\n",
    "\n",
    "    #torch.save(unet.to('cpu'), 'saves/unet.model')\n",
    "\n",
    "\n",
    "train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "22f955b2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a72f7737ddf24b868be27b18c5b694a2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Upload 1 LFS files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "711d4f375af84ff59ef05920c646f3bf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/3.44G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/lansinuote/diffsion_from_scratch.unet/commit/32f5e4163edb6d1a3fa1d8265ad2cdf0406cb425', commit_message='Upload model', commit_description='', oid='32f5e4163edb6d1a3fa1d8265ad2cdf0406cb425', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import PreTrainedModel, PretrainedConfig\n",
    "\n",
    "\n",
    "#包装类\n",
    "class Model(PreTrainedModel):\n",
    "    config_class = PretrainedConfig\n",
    "\n",
    "    def __init__(self, config):\n",
    "        super().__init__(config)\n",
    "        self.unet = unet.to('cpu')\n",
    "\n",
    "#保存到hub\n",
    "Model(PretrainedConfig()).push_to_hub(\n",
    "    repo_id='lansinuote/diffsion_from_scratch.unet',\n",
    "    use_auth_token=open('/root/hub_token.txt').read().strip())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pt39]",
   "language": "python",
   "name": "conda-env-pt39-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
